
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg') 

min_length = 9999

# 显示图像的辅助函数
def show_img(title, img, is_gray=False, save_path=None):
    plt.figure(figsize=(8, 6))
    if is_gray:
        plt.imshow(img, cmap='gray')
    else:
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.title(title)
    plt.axis('off')
    
    # 只保存图片，不显示
    if save_path:
        cv2.imwrite(save_path, img)
        print(f"标记角点的图像已保存至：{save_path}")
    
    # 关闭当前画布，释放内存（避免内存占用过高）
    plt.close()

# 1. 提取白纸区域
def extract_white_paper(img):
    # 转换到HSV颜色空间（便于分离亮度和饱和度）
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    
    # 白纸特征：低饱和度（s）+ 高亮度（v）
    # 阈值可根据实际环境调整（此处为示例值）
    s_threshold = 60  # 饱和度阈值（低饱和度）
    v_threshold = 180  # 亮度阈值（高亮度）
    s_mask = s < s_threshold
    v_mask = v > v_threshold
    white_mask = np.logical_and(s_mask, v_mask).astype(np.uint8) * 255  # 合并掩码
    
    # 形态学操作：去除小噪声，填充孔洞
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_CLOSE, kernel, iterations=3)  # 填充孔洞
    white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_OPEN, kernel, iterations=1)   # 去除小噪声
    
    # 查找轮廓，筛选最大的连通区域（假设白纸是最大的白色区域）
    contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None  # 未找到白纸区域
    max_contour = max(contours, key=cv2.contourArea)
    area = cv2.contourArea(max_contour)
    if area < 10000:  # 过滤过小的区域（避免误判）
        return None, None
    
    # 创建白纸区域的掩码（仅保留最大轮廓内的区域）
    paper_mask = np.zeros_like(white_mask)
    cv2.drawContours(paper_mask, [max_contour], -1, 255, -1)
    
    # 提取白纸区域（原图中仅保留掩码内的像素）
    paper_img = cv2.bitwise_and(img, img, mask=paper_mask)
    return paper_img, paper_mask

# 2. 从白纸中提取线条图形（假设线条为深色）
def extract_lines(paper_img, paper_mask):
    # 转换为灰度图
    gray = cv2.cvtColor(paper_img, cv2.COLOR_BGR2GRAY)
    
    # 反色处理（让线条变为白色，背景为黑色）
    # 先通过掩码去除白纸外的干扰（置为黑色）
    gray[paper_mask == 0] = 0
    
    inverted = cv2.bitwise_not(gray)  # 反色：白纸（亮）→ 黑，线条（暗）→ 白
    '''
    plt.figure()
    plt.imshow(inverted, cmap='gray')
    plt.title("步骤3：反色处理（线条为白，背景为黑）")
    plt.axis('off')
    plt.show()
    '''
    # 通过掩码去除白纸外的干扰（置为黑色）
    inverted[paper_mask == 0] = 0
    '''
    plt.figure()
    plt.imshow(inverted, cmap='gray')
    plt.title("步骤3：反色处理（线条为白，背景为黑）")
    plt.axis('off')
    plt.show()
    '''
    # 阈值处理：保留线条（白色）
    _, line_mask = cv2.threshold(inverted, 80, 255, cv2.THRESH_BINARY)  # 阈值可调整
    #line_mask[line_mask!=0]=255
    '''
    plt.figure()
    plt.imshow(line_mask, cmap='gray')
    plt.title("步骤4：阈值处理（提取线条）")
    plt.axis('off')
    plt.show()
    '''
    # 形态学优化：连接断裂的线条，去除小噪声
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_CLOSE, kernel, iterations=2)  # 连接线条
    line_mask = cv2.morphologyEx(line_mask, cv2.MORPH_OPEN, kernel, iterations=1)   # 去噪声
    return line_mask

def prune_short_branches(skeleton, min_length):
    img = skeleton.copy()
    _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
    pruned = binary.copy()
    visited = np.zeros_like(binary, dtype=bool)
    devisited = np.zeros_like(binary, dtype=bool)
    # 8邻域方向
    neighbors = [(-1, -1), (-1, 0), (-1, 1),
                 (0, -1),          (0, 1),
                 (1, -1),  (1, 0), (1, 1)]
    def be_endpoint(img):
        for x in range(binary.shape[0]):
            for y in range(binary.shape[1]):
                if is_endpoint(x, y, img)==True:
                    return True
        return False

    def is_endpoint(x, y,img):
        if img[x, y] != 255:
            return False
        count = 0
        for dx, dy in neighbors:
            nx, ny = x + dx, y + dy
            if 0 <= nx < img.shape[0] and 0 <= ny < img.shape[1]:
                if img[nx, ny] == 255:
                    count += 1
        return count == 1
    
    def dfs_branch_iterative(start_x, start_y,img):
        # 栈元素：(当前x, 当前y, 上一x, 上一y)
        stack = [(start_x, start_y, -1, -1)]
        path = []
        visited[start_x, start_y] = True  # 标记起点为已访问
        
        while stack:
            x, y, prev_x, prev_y = stack.pop()  # 弹出栈顶元素处理
            path.append((x, y))  # 记录当前像素到路径
            
            # 寻找下一个有效像素（排除上一个像素）
            next_pixels = []
            for dx, dy in neighbors:
                nx, ny = x + dx, y + dy
                if 0 <= nx < img.shape[0] and 0 <= ny < img.shape[1]:
                    # 条件：前景像素、未访问、不是上一个像素
                    if img[nx, ny] == 255 and ((not visited[nx, ny]) or devisited[nx, ny]) and (nx != prev_x or ny != prev_y):
                        next_pixels.append((nx, ny))
            
            # 分支点：若有多个下一跳，当前分支结束（只沿一个方向继续遍历）
            if len(next_pixels) == 1:
                nx, ny = next_pixels[0]
                visited[nx, ny] = True  # 提前标记为已访问，避免重复处理
                stack.append((nx, ny, x, y))  # 下一个像素入栈
            # 若有0个下一跳（终点）或≥2个下一跳（分支点），则当前分支结束
            else:
                break
        
        return path
    
    # 遍历图像，检测所有端点并处理
    while be_endpoint(pruned):
        for x in range(binary.shape[0]):
            for y in range(binary.shape[1]):
            # 仅处理未访问的端点
                if binary[x, y] == 255 and ((not visited[x, y]) or devisited[x, y]) and is_endpoint(x, y,pruned):
                # 用迭代版DFS替代递归，避免深度问题
                    branch_path = dfs_branch_iterative(x, y,pruned)
                    branch_length = len(branch_path)

                # 显示中间结果（可根据需要保留或删除）
                    ''' 
                    cv2.namedWindow("skeleton", cv2.WINDOW_NORMAL)
                    cv2.imshow("skeleton", pruned)
                    cv2.waitKey(0)
                    cv2.destroyAllWindows()
                    '''
                # 若分支长度小于阈值，则删除该分支（忽略最后一个元素）
                    if branch_length < min_length:
                        devisited[branch_path[-1]] = True
                        for (px, py) in branch_path[:-1]:  # 忽略最后一个元素
                            pruned[px, py] = 0

    
    return pruned  
  

def skeletonize(line_mask):
    skeleton = cv2.ximgproc.thinning(line_mask,thinningType=cv2.ximgproc.THINNING_GUOHALL)
    
    #skeleton = cv2.morphologyEx(skeleton,cv2.MORPH_CLOSE,kernel,iterations=1)
    cv2.imwrite("OutPut/unprunned.png", skeleton)
    skeleton = prune_short_branches(skeleton,min_length)
    
    return skeleton

# 主函数
def main(img_path):
    # 读取图像
    img = cv2.imread(img_path)
    if img is None:
        print("无法读取图像，请检查路径")
        return
    # 保存原图（可选）
    show_img("原图", img, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/original.png")
    
    # 步骤1：提取白纸区域
    paper_img, paper_mask = extract_white_paper(img)
    if paper_img is None:
        print("未检测到白纸区域")
        return
    show_img("提取的白纸区域", paper_img, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/white_paper.png")
    show_img("白纸掩码", paper_mask, is_gray=True, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/paper_mask.png")
    
    # 步骤2：提取线条图形
    line_mask = extract_lines(paper_img, paper_mask)
    show_img("提取的线条", line_mask, is_gray=True, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/line_mask.png")
    # 步骤3：生成骨架
    skeleton = skeletonize(line_mask)

    show_img("线条骨架（中轴线）", skeleton, is_gray=True, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/skeleton.png")
    #detect_skeleton_corners(skeleton,"/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/Corner.png")
    #detect_skeleton_features(skeleton,"/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/Features.png")
    temp_morph = line_mask.copy()
    temp_mask = (skeleton==255)
    temp_morph[temp_mask]= 0
    show_img("1", temp_morph, is_gray=True, save_path="/home/gerinskinut/Project/Dreamc/Dreamc/OutPut/result_together.png")
    
if __name__ == "__main__":
    main("InPut/misty--.jpg")  # 此处替换为实际图像路径